import torch

x = [1.0, 2.0, 3.0]
y = [6.0, 17.0, 34.0]
w1 = torch.Tensor([1.0])
w1.requires_grad = True
w2 = torch.Tensor([1.0])
w2.requires_grad = True
b = torch.Tensor([2.0])
b.requires_grad = True


def forward(x):
    return x * x * w1 + x * w2 + b


def loss(x, y):
    y_pred = forward(x)
    return (y_pred - y) ** 2


for epoch in range(100):
    for x_data, y_data in zip(x, y):
        l = loss(x_data, y_data)
        l.backward()
        w1.data = w1.data - 0.01 * w1.grad.data
        w2.data = w2.data - 0.01 * w2.grad.data
        b.data = b.data - 0.01 * b.grad.data
        w1.grad.data.zero_()
        w2.grad.data.zero_()
        b.grad.data.zero_()

print(forward(4))
print(w1, w2, b)
